#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
拦截重放服务
"""

import logging
import json
import requests
from datetime import datetime
from bson import ObjectId
from db.mongo import get_db
from intercept.models.intercept import InterceptServiceModel, InterceptRequestRecordModel

logger = logging.getLogger(__name__)

class InterceptService:
    """拦截重放服务"""

    @staticmethod
    def create_intercept_service(service_name, service_ip=None, service_port=None):
        """
        创建拦截服务

        Args:
            service_name: 服务名称
            service_ip: 服务IP地址
            service_port: 服务端口

        Returns:
            str: 服务ID
        """
        db = get_db()

        # 检查服务名称是否已存在
        existing_service = db.intercept_services.find_one({"service_name": service_name})
        if existing_service:
            return existing_service["_id"]

        # 创建拦截服务
        service = InterceptServiceModel(
            service_name=service_name,
            service_ip=service_ip,
            service_port=service_port
        )
        service_dict = service.to_dict()

        # 插入拦截服务
        db.intercept_services.insert_one(service_dict)
        service_id = service_dict["_id"]

        return service_id

    @staticmethod
    def get_intercept_services(page=1, page_size=10):
        """
        获取拦截服务列表

        Args:
            page: 页码
            page_size: 每页大小

        Returns:
            dict: 包含拦截服务列表和总数的字典
        """
        db = get_db()

        # 计算跳过的记录数
        skip = (page - 1) * page_size

        # 查询拦截服务
        services = list(db.intercept_services.find().sort("created_at", -1).skip(skip).limit(page_size))

        # 获取总记录数
        total = db.intercept_services.count_documents({})

        return {
            "services": services,
            "total": total,
            "page": page,
            "page_size": page_size,
            "pages": (total + page_size - 1) // page_size
        }

    @staticmethod
    def get_intercept_service_by_id(service_id):
        """
        根据ID获取拦截服务

        Args:
            service_id: 服务ID

        Returns:
            dict: 拦截服务详情
        """
        db = get_db()

        # 将 service_id 转换为 ObjectId
        if isinstance(service_id, str):
            try:
                service_id = ObjectId(service_id)
            except Exception as e:
                logger.error(f"无效的 service_id: {service_id}, 错误: {e}")
                return None

        # 查询拦截服务
        # 先尝试使用 ObjectId 查询
        service = db.intercept_services.find_one({"_id": service_id})

        # 如果没有找到，尝试使用字符串查询
        if not service and isinstance(service_id, ObjectId):
            service = db.intercept_services.find_one({"_id": str(service_id)})

        # 如果还是没有找到，尝试直接使用原始的 service_id 查询
        if not service and isinstance(service_id, ObjectId):
            original_service_id = service_id
            service = db.intercept_services.find_one({"_id": original_service_id})
        if not service:
            return None

        # 将 ObjectId 转换为字符串
        service["_id"] = str(service["_id"])

        return service

    @staticmethod
    def get_intercept_service_by_name(service_name):
        """
        根据名称获取拦截服务

        Args:
            service_name: 服务名称

        Returns:
            dict: 拦截服务详情
        """
        db = get_db()

        # 查询拦截服务
        service = db.intercept_services.find_one({"service_name": service_name})
        if not service:
            return None

        return service

    @staticmethod
    def update_intercept_service(service_id, service_name=None, service_ip=None, service_port=None, is_active=None):
        """
        更新拦截服务

        Args:
            service_id: 服务ID
            service_name: 服务名称
            service_ip: 服务IP地址
            service_port: 服务端口
            is_active: 是否激活

        Returns:
            bool: 是否成功更新
        """
        db = get_db()

        # 查询拦截服务
        service = db.intercept_services.find_one({"_id": service_id})
        if not service:
            return False

        # 准备更新数据
        update_data = {"updated_at": datetime.utcnow()}
        if service_name is not None:
            update_data["service_name"] = service_name
        if service_ip is not None:
            update_data["service_ip"] = service_ip
        if service_port is not None:
            update_data["service_port"] = service_port
        if is_active is not None:
            update_data["is_active"] = is_active

        # 更新拦截服务
        db.intercept_services.update_one({"_id": service_id}, {"$set": update_data})

        return True

    @staticmethod
    def delete_intercept_service(service_id):
        """
        删除拦截服务

        Args:
            service_id: 服务ID

        Returns:
            bool: 是否成功删除
        """
        db = get_db()

        # 查询拦截服务
        service = db.intercept_services.find_one({"_id": service_id})
        if not service:
            return False

        # 删除相关的请求记录
        db.intercept_request_records.delete_many({"intercept_service_id": service_id})

        # 删除拦截服务
        db.intercept_services.delete_one({"_id": service_id})

        return True

    @staticmethod
    def find_matching_request_record(service_id, method, url):
        """
        查找匹配的请求记录

        Args:
            service_id: 服务ID
            method: 请求方法
            url: 请求URL

        Returns:
            dict: 匹配的请求记录
        """
        db = get_db()

        # 查询匹配的请求记录
        record = db.intercept_request_records.find_one({
            "intercept_service_id": service_id,
            "request_method": method,
            "request_url": url
        })

        return record

    @staticmethod
    def create_or_update_request_record(service_id, method, url, request_headers, request_body,
                                       response_status, response_headers, response_body):
        """
        创建或更新请求记录

        Args:
            service_id: 服务ID
            method: 请求方法
            url: 请求URL
            request_headers: 请求头
            request_body: 请求体
            response_status: 响应状态码
            response_headers: 响应头
            response_body: 响应体

        Returns:
            str: 请求记录ID
        """
        db = get_db()

        # 查找匹配的请求记录
        existing_record = InterceptService.find_matching_request_record(service_id, method, url)

        if existing_record:
            # 更新现有记录
            record_id = existing_record["_id"]
            db.intercept_request_records.update_one(
                {"_id": record_id},
                {"$set": {
                    "request_headers": request_headers,
                    "request_body": request_body,
                    "response_status": response_status,
                    "response_headers": response_headers,
                    "response_body": response_body,
                    "updated_at": datetime.utcnow()
                }}
            )
            return record_id
        else:
            # 创建新记录
            record = InterceptRequestRecordModel(
                intercept_service_id=service_id,
                request_method=method,
                request_url=url,
                request_headers=request_headers,
                request_body=request_body,
                response_status=response_status,
                response_headers=response_headers,
                response_body=response_body
            )
            record_dict = record.to_dict()

            # 插入请求记录
            db.intercept_request_records.insert_one(record_dict)
            record_id = record_dict["_id"]

            return record_id

    @staticmethod
    def update_replay_count(record_id):
        """
        更新重放计数

        Args:
            record_id: 请求记录ID

        Returns:
            bool: 是否成功更新
        """
        db = get_db()

        # 查询请求记录
        record = db.intercept_request_records.find_one({"_id": record_id})
        if not record:
            return False

        # 更新重放计数
        db.intercept_request_records.update_one(
            {"_id": record_id},
            {"$inc": {"replay_count": 1}, "$set": {"last_replayed_at": datetime.utcnow()}}
        )

        return True

    @staticmethod
    def get_request_records(service_id=None, page=1, page_size=10):
        """
        获取请求记录列表

        Args:
            service_id: 服务ID（可选）
            page: 页码
            page_size: 每页大小

        Returns:
            dict: 包含请求记录列表和总数的字典
        """
        db = get_db()

        # 构建查询条件
        query = {}
        if service_id:
            query["intercept_service_id"] = service_id

        # 计算跳过的记录数
        skip = (page - 1) * page_size

        # 查询请求记录
        records = list(db.intercept_request_records.find(query).sort("created_at", -1).skip(skip).limit(page_size))

        # 获取总记录数
        total = db.intercept_request_records.count_documents(query)

        return {
            "records": records,
            "total": total,
            "page": page,
            "page_size": page_size,
            "pages": (total + page_size - 1) // page_size
        }

    @staticmethod
    def get_request_record_by_id(record_id):
        """
        根据ID获取请求记录

        Args:
            record_id: 请求记录ID

        Returns:
            dict: 请求记录详情
        """
        db = get_db()

        # 查询请求记录
        record = db.intercept_request_records.find_one({"_id": record_id})
        if not record:
            return None

        return record

    @staticmethod
    def delete_request_record(record_id):
        """
        删除请求记录

        Args:
            record_id: 请求记录ID

        Returns:
            bool: 是否成功删除
        """
        db = get_db()

        # 查询请求记录
        record = db.intercept_request_records.find_one({"_id": record_id})
        if not record:
            return False

        # 删除请求记录
        db.intercept_request_records.delete_one({"_id": record_id})

        return True

    @staticmethod
    def intercept_request(service_name, method, url, headers, body):
        """
        拦截请求

        Args:
            service_name: 服务名称
            method: 请求方法
            url: 请求URL
            headers: 请求头
            body: 请求体

        Returns:
            tuple: (是否拦截, 响应状态码, 响应头, 响应体, 请求记录ID)
        """
        # 获取拦截服务
        service = InterceptService.get_intercept_service_by_name(service_name)
        if not service or not service.get("is_active", False):
            return False, None, None, None, None

        service_id = service["_id"]

        # 查找匹配的请求记录
        record = InterceptService.find_matching_request_record(service_id, method, url)
        if record:
            # 更新重放计数
            InterceptService.update_replay_count(record["_id"])

            # 返回记录的响应
            response_headers = record.get("response_headers", {})
            response_headers["X-Replayed-Response"] = "true"
            response_headers["X-Replay-Count"] = str(record.get("replay_count", 0) + 1)
            if record.get("created_at"):
                response_headers["X-Original-Recorded-At"] = record["created_at"].isoformat()

            return True, record.get("response_status"), response_headers, record.get("response_body"), record["_id"]
        else:
            # 转发请求并记录
            try:
                # 发送请求
                response = requests.request(
                    method=method,
                    url=url,
                    headers=headers,
                    data=body,
                    allow_redirects=False,
                    timeout=30
                )

                # 记录请求和响应
                record_id = InterceptService.create_or_update_request_record(
                    service_id,
                    method,
                    url,
                    headers,
                    body,
                    response.status_code,
                    dict(response.headers),
                    response.text
                )

                # 返回原始响应
                return False, response.status_code, dict(response.headers), response.text, record_id
            except Exception as e:
                logger.error(f"转发请求失败: {e}")
                return False, 500, {"Content-Type": "application/json"}, json.dumps({"error": str(e)}), None
